/*
 *  btree.c
 *  rr-module-text
 *
 *  Created by Rodrigo Sestari on 02/04/11.
 *  Copyright 2011 __MyCompanyName__. All rights reserved.
 *
 */


#include <stdlib.h>
#include <stdio.h>
#include "btree.h"

#define BALANCED        0
#define LEFT_HEAVY      1
#define RIGHT_HEAVY     -1

struct AVLnode
{
    void *key;
    void *data;
    int balance;
    int hidden;
    struct AVLnode *left;
    struct AVLnode *right;
};

static int insert(struct btree *tree, struct AVLnode **t,
				  int *balanced, void *key, void *data);
static int rotate_right(struct AVLnode **t);
static int rotate_left(struct AVLnode **t);
static void destroy(struct btree *tree, struct AVLnode *t);
static int remove_node(struct btree *tree, struct AVLnode *t, void *key);
static void *lookup(struct btree *tree, struct AVLnode *t, void *key);
static void inorder(struct AVLnode *n, void *topass,
                    void (*print)(void *key, void *data, void *topass));

struct btree *btree_create(void (*destroy_data)(void *),
                           void (*destroy_key)(void *),
                           int (*compare)(void *, void *))
{
    struct btree *t = malloc(sizeof(struct btree));
    t->head = NULL;
    t->size = 0;
    t->destroy_data = destroy_data;
    t->destroy_key = destroy_key;
    t->compare = compare;
	
    return t;
}

void btree_destroy(struct btree *tree)
{
    destroy(tree, tree->head);
    free(tree);
}

int btree_insert(struct btree *tree, void *key, void *data)
{
    int balanced = 0;
    int retval = insert(tree, &(tree->head), &balanced, key, data);
    if (!retval)
        tree->size++;
    return retval;
}

int btree_remove(struct btree *tree, void *key)
{
    return remove_node(tree, tree->head, key);
}

void *btree_lookup(struct btree *tree, void *key)
{
    return lookup(tree, tree->head, key);
}

int btree_size(struct btree *tree)
{
    return tree->size;
}

void btree_dump(struct btree *tree, void *topass,
                void (*print)(void *key, void *data, void *topass))
{
    inorder(tree->head, topass, print);
}


static int insert(struct btree *tree, struct AVLnode **t, int *balanced,
				  void *key, void *data)
{
    int comp, retval;
    struct AVLnode *node;
	
    if (!(*t)) {
        node = malloc(sizeof(struct AVLnode));
        node->key = key;
        node->data = data;
        node->balance = BALANCED;
        node->hidden = 0;
        node->left = NULL;
        node->right = NULL;
        *t = node;
        return 0;
    }
    else {
        comp = tree->compare(key, (*t)->key);
        if (comp < 0) {   /* Add to left */
            if (!(*t)->left) {
                node = malloc(sizeof(struct AVLnode));
                node->key = key;
                node->data = data;
                node->balance = BALANCED;
                node->hidden = 0;
                node->left = NULL;
                node->right = NULL;
                (*t)->left = node;
                *balanced = 0;
            }
            else {
                retval = insert(tree, &((*t)->left), balanced, key, data);
                if (retval != 0)
                    return retval;
            }
            if (!(*balanced)) {
                switch ((*t)->balance) {
                    case LEFT_HEAVY:
                        rotate_left(t);
                        *balanced = 1;
                        break;
                    case BALANCED:
                        (*t)->balance = LEFT_HEAVY;
                        break;
                    case RIGHT_HEAVY:
                        (*t)->balance = BALANCED;
                        *balanced = 1;
                        break;
                }
            }
        }
        else if (comp > 0) {  /* try the right */
            if (!(*t)->right) {
                node = malloc(sizeof(struct AVLnode));
                node->key = key;
                node->data = data;
                node->balance = BALANCED;
                node->hidden = 0;
                node->left = NULL;
                node->right = NULL;
                (*t)->right = node;
                *balanced = 0;
            }
            else {
                retval = insert(tree, &((*t)->right), balanced, key, data);
                if (retval)
                    return retval;
            }
            if (!(*balanced)) {
                switch ((*t)->balance) {
                    case LEFT_HEAVY:
                        (*t)->balance = BALANCED;
                        *balanced = 1;
                        break;
                    case BALANCED:
                        (*t)->balance = RIGHT_HEAVY;
                        break;
                    case RIGHT_HEAVY:
                        rotate_right(t);
                        *balanced = 1;
                        break;
                }
            }
        }
        else { /* We've come upon data previously hidden */
            if ((*t)->hidden == 0)
                return 1;
            else {
                if (tree->destroy_data)
                    tree->destroy_data((*t)->data);
                (*t)->data = data;
                (*t)->hidden = 0;
                *balanced = 1;
            }
        }
    }
    return 0;
}

static int rotate_right(struct AVLnode **t)
{
    struct AVLnode *ri = (*t)->right;
    struct AVLnode *gchild;
	
    if (ri->balance == RIGHT_HEAVY) {
        /* RR rotation */
        (*t)->right = ri->left;
        ri->left = *t;
        (*t)->balance = BALANCED;
        ri->balance = BALANCED;
        *t = ri;
    }
    else {
        /* RL rotation */
        gchild = ri->left;
        ri->left = gchild->right;
        gchild->right = ri;
        (*t)->right = gchild->left;
        gchild->left = *t;
		
        switch (gchild->balance) {
            case LEFT_HEAVY:
                (*t)->balance = BALANCED;
                ri->balance = RIGHT_HEAVY;
                break;
            case BALANCED:
                (*t)->balance = BALANCED;
                ri->balance = BALANCED;
                break;
            case RIGHT_HEAVY:
                (*t)->balance = LEFT_HEAVY;
                ri->balance = BALANCED;
                break;
        }
        gchild->balance = BALANCED;
        *t = gchild;
    }
    return 0;
}


static int rotate_left(struct AVLnode **t)
{
    struct AVLnode *left = (*t)->left;
    struct AVLnode *gchild;
	
    if (left->balance == LEFT_HEAVY) {
        /* LL rotation */
        (*t)->left = left->right;
        left->right = *t;
        (*t)->balance = BALANCED;
        left->balance = BALANCED;
        *t = left;
    }
    else {
        /* LR rotation */
        gchild = left->right;
        left->right = gchild->left;
        gchild->left = left;
        (*t)->left = gchild->right;
        gchild->right = *t;
		
        switch (gchild->balance) {
            case LEFT_HEAVY:
                (*t)->balance = RIGHT_HEAVY;
                left->balance = BALANCED;
                break;
            case BALANCED:
                (*t)->balance = BALANCED;
                left->balance = BALANCED;
                break;
            case RIGHT_HEAVY:
                (*t)->balance = BALANCED;
                left->balance = LEFT_HEAVY;
                break;
        }
        gchild->balance = BALANCED;
        *t = gchild;
    }
    return 0;
}

static int remove_node(struct btree *tree, struct AVLnode *t, void *key)
{
    int comp;
    if (t) {
        comp = tree->compare(key, t->key);
        if (comp == 0) {
            t->hidden = 1;
            tree->size--;
            return 0;
        }
        else if (comp < 0)
            return remove_node(tree, t->left, key);
        else if (comp > 0)
            return remove_node(tree, t->right, key);
    }
    return 1;
}

static void destroy(struct btree *tree, struct AVLnode *t)
{
    if (t) {
        destroy(tree, t->left);
        destroy(tree, t->right);
        if (tree->destroy_key)
            tree->destroy_key(t->key);
        if (tree->destroy_data)
            tree->destroy_data(t->data);
        free(t);
    }
}

static void *lookup(struct btree *tree, struct AVLnode *t, void *key)
{
    int comp;
	
    if (t) {
        comp = tree->compare(key, t->key);
        if (comp < 0)
            return lookup(tree, t->left, key);
        else if (comp > 0)
            return lookup(tree, t->right, key);
        else if (!t->hidden)
            return t->data;
        else
            return NULL;
    }
    else
        return NULL;
}

static void inorder(struct AVLnode *n, void *topass,
                    void (*print)(void *key, void *data, void *topass))
{
    if (n) {
        inorder(n->left, topass, print);
        if (!(n->hidden)) {
            /* printf("Balance: %d\n", n->balance); */
            print(n->key, n->data, topass);
        }
        inorder(n->right, topass, print);
    }
}
